package algorithmic_data_structure.堆与堆排序;

import org.junit.Test;

import java.util.Comparator;
import java.util.PriorityQueue;

public class P502_IPO {
    @Test
    public void test(){
        System.out.println(findMaximizedCapital(2, 0, new int[]{1, 2, 3}, new int[]{0, 1, 1}));
        System.out.println(findMaximizedCapital(10, 0, new int[]{1, 2, 3}, new int[]{0, 1, 2}));
    }
    public int findMaximizedCapital(int k, int w, int[] profits, int[] capital) {
        PriorityQueue<Integer> minCapitalHeap = new PriorityQueue<>(new Comparator<Integer>() {
            @Override
            public int compare(Integer index1, Integer index2) {
                return capital[index1] - capital[index2];
            }
        });
        PriorityQueue<Integer> maxProfitsHeap = new PriorityQueue<>(
                (i1,i2) ->{return profits[i2] - profits[i1];});
        for (int i = 0; i < capital.length; i++) {
            minCapitalHeap.add(i);
        }
        int maxP = w;
        for (int i = 0; i < k; i++) {
            while (!minCapitalHeap.isEmpty() && maxP >= capital[minCapitalHeap.peek()]){
                maxProfitsHeap.add(minCapitalHeap.poll());
            }
            if (maxProfitsHeap.isEmpty()){
                break;
            }
            Integer choose = maxProfitsHeap.poll();
            maxP += profits[choose];
        }
        return maxP;
    }
}
